"""Simple MNIST image classifier example with LightningModule and LightningDataModule.

To run: python minmst_example.py 
"""

import torch
import torch_npu
import os

from examples.dataset_moudle.mnist_datamodule import MNISTDataModule
from examples.module.mnist_module import ImageClassifier
from pytorch_lightning.utilities.cli import LightningCLI
from pytorch_lightning.loggers import CometLogger,NeptuneLogger,TensorBoardLogger,WandbLogger,MLFlowLogger

from lightning_npu.accelerators.npu import NPUAccelerator
from lightning_npu.strategies.npu_parallel import NPUParallelStrategy


def cli_main():
    # The LightningCLI removes all the boilerplate associated with arguments parsing. This is purely optional.
    comet_logger = CometLogger(
        api_key="comet_API_KEY",
        save_dir=".",
        project_name="MNIST",
        rest_api_key="comet_API_KEY",
        experiment_name="comet_logs",
    )
    mlf_logger = MLFlowLogger(experiment_name="mlf_logs", tracking_uri="to_path")
    neptune_logger = NeptuneLogger(
        api_key="neptune_API_KEY",
        project="WORKSPACE/PROJECT",
        tags=["training","minist"]
    )
    tb_logger = TensorBoardLogger("to_path", name="tb_logs")
    wandb_logger = WandbLogger(project="MNIST", log_model="all")
    cli = LightningCLI(
        ImageClassifier,
        MNISTDataModule,
        trainer_defaults={
            "accelerator": NPUAccelerator(),
            "devices": [0,1,2,3],
            "max_epochs": 10,
            "strategy": NPUParallelStrategy(),
            "logger": [comet_logger,neptune_logger,tb_logger,wandb_logger,mlf_logger]
        },
        seed_everything_default=42,
        save_config_overwrite=True,
        run=False
    )
    wandb_logger.watch(ImageClassifier())
    cli.trainer.fit(cli.model, datamodule=cli.datamodule)
    cli.trainer.test(ckpt_path="best", datamodule=cli.datamodule)




if __name__ == "__main__":
    cli_main()